import scipy.ndimage as ndi
import numpy as np
import os
import netCDF4 as nc
import gc

def cal_nn(lbl):
    lblmax = lbl.max()
    nn = []
    for i in range(1,lblmax+1):
        nn.append(np.sum(np.where(lbl==i,1,0)))
    return nn


def get_idx_size(lbl,idx):
    tmp = np.where(lbl == idx, 1, 0)
    idxsize = np.sum(tmp)
    return idxsize

def filter(var,varcrit,ncrit):
    s=np.array([[1,1,1],
                [1,1,1],
                [1,1,1]])
    var=np.where((np.isnan(var))|(var>100),0,var)
    zhi0=np.where(var>varcrit,1,0)
    lbl = ndi.label(zhi0,structure=s)[0]
    lblmax = lbl.max()
    tmp = np.zeros(lbl.shape, 'int')
    for i in range(1, lblmax + 1):
        idxsize=get_idx_size(lbl,i)
        tmp = np.where(lbl == i, idxsize, tmp)
    newvar = np.where(tmp > ncrit, 1, 0)  # 太小的就不要了
    newlbl = ndi.label(newvar,structure=s)[0]

    return newlbl


def strinarray(str,var):
    ny=var.shape[0]
    out = np.zeros(var.shape, 'str')
    for j in range(ny):
        ceshi = np.array([str in x for x in var[j, :]])
        out[j, :] = ceshi
    return out

def getmaxlabel(var):
    var0=[int(x) for x in var.flatten() ]
    varmax=max(var0)
    return varmax

def getelements(var):
    out=set(var.flatten())
    out.remove(0)
    return out
def overlap(l0,l1,l2,maxini,info_file,time):
    l0max=l0.max()
    l1=-1*l1
    l1min=l1.min()
    for i1 in np.arange(l1min,0): # 对于旧label中有，新label中没有的，以后就不用这个编号了
        flag=0
        tmp1=np.where(l1==i1,1,0)
        tmplabel=[]
        for i0 in range(1,l0max+1):
            tmp0=np.where(l0==i0,1,0)  ### 上一个时刻的label中 等于某一个label的位置的地方=1，否则等于0
            tmpsum=np.sum(tmp0*tmp1)  # 重叠范围内有多少个点
            if tmpsum>=5:
                tmplabel.append(i0)

                flag=1
        if flag==1:
            if len(tmplabel)==1:
                l1=np.where(l1==i1,tmplabel[0],l1)  ### 只有一个编号重叠，那就简单地继承这个编号
            elif len(tmplabel)>1:
                sizes=[get_idx_size(l0,zhi) for zhi in tmplabel]
                idx_maxsize=np.argmax(np.array(sizes))
                label_maxsize=tmplabel[idx_maxsize]
                l1=np.where(l1==i1,label_maxsize,l1)   ### 如果一个对流由多个对流合并而来，那它继承面积最大的那个对流的编号
                for idx_tmp,lbl_tmp in enumerate(tmplabel):
                    if idx_tmp!=idx_maxsize:
                        info_txt=time+' '+str(lbl_tmp)+' merge into '+str(label_maxsize)+'\n'
                        info_file.writelines(info_txt)
        if flag==0: # 没有重叠,看情况是否给新的编号
            pos1=np.where(l1==i1)  # 等于i1的地方
            l2sub=l2[pos1]
            l2tmp=np.where(l2sub>0,1,0)
            l1_l2_ovlap=np.sum(l2tmp) # l2的元素 中 有没有和l1=i1 的地方有重合
            if l1_l2_ovlap>=5: ## 只有生命史至少大于1,而且重叠区域大于5 的 才会给编号
                maxini = maxini + 1
                l1=np.where(l1==i1,maxini,l1)
            else:
                l1=np.where(l1==i1,0,l1)

    # del(tmp1);del(tmp0);del(tmpsum);del(l2sub);del(l2tmp)

    s = np.array([[1, 1, 1],
                  [1, 1, 1],
                  [1, 1, 1]])
    newl1ele=getelements(l1)

    for ele in newl1ele:
        tmp=np.where(l1==ele,1,0)
        lbltmp = ndi.label(tmp,structure=s)[0]
        tmpmax = lbltmp.max()
        if tmpmax>1:  ## 表示有多个单体 共用一个label，那么他们现在label名字变为：继承的label+'-'+序号
            idxs = range(1, tmpmax + 1)
            sizes = [get_idx_size(lbltmp, zhi) for zhi in idxs]
            idx_maxsize = np.argmax(np.array(sizes))
            for idx in idxs:
                if idx==idx_maxsize+1:
                    continue
                else:
                    label_size=np.where(lbltmp==idx)[0].shape[0]
                    if label_size>=5:
                        maxini = maxini + 1
                        l1 = np.where(lbltmp == idx, maxini, l1)
                        info_txt=time+' '+str(ele)+' split into '+str(maxini)+'\n'
                        info_file.writelines(info_txt)
                    else:
                        l1 = np.where(lbltmp == idx, 0, l1)

    return l1




date='20250701'

### merge and split info file

infodir=r'E:\akesu_radar\cheng\20250701_info'
infofilename=os.path.join(infodir,date+'_split_merge_info.txt')
info_file=open(infofilename,'w')

fdir=os.path.join(r'E:\akesu_radar\cheng',date)
files=sorted([os.path.join(fdir,x) for x in os.listdir(fdir) if 'npz' in x])
outdir=r'E:\akesu_radar\cheng\\'+date+'_label'
if not os.path.exists(outdir):
    os.makedirs(outdir)
nt=len(files)

varcrit=35
ncrit=5

# ### ini file###

f0=np.load(files[0])
cr0=f0['CR']
lbl0=filter(cr0,varcrit,ncrit)
maxini=getmaxlabel(lbl0)
time=os.path.basename(files[0]).split('_')[4]
np.save(os.path.join(outdir,time+'_label'),lbl0)

outfiles=os.listdir(outdir)
outtimes=[x.split('_')[0] for x in outfiles]

for it in range(1,nt-1):
    time = os.path.basename(files[it]).split('_')[4]
    if time in outtimes:
        continue
    print(time)
    f1=np.load(files[it])
    cr1=f1['CR']
    lbl1=filter(cr1,varcrit,ncrit)

    f2 = np.load(files[it+1])
    cr2 = f2['CR']
    lbl2 = filter(cr2, varcrit, ncrit)


    sp1=cr1.shape
    sp2=cr2.shape
    print(sp1,sp2)
    # if sp1!=(555,593) or sp1!=sp2:
    #     print(time,sp1,sp2)


    newlbl1=overlap(lbl0,lbl1,lbl2,maxini,info_file,time)
    np.save(os.path.join(outdir,time+'_label'),newlbl1)
    maxini=max(maxini,newlbl1.max())
    lbl0=newlbl1
    del(newlbl1);del(cr1);del(lbl1);del(f1)
    del(cr2);del(lbl2);del(f2)
    gc.collect()

info_file.close()








